import threading
from math import *
import time
import sys
from scipy.optimize import root

import rclpy
from rclpy.action import ActionClient
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor

from my_action.action import JointTarget
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener
import tf_transformations
import os

def feedback_cb(logger, joint_type, feedback):
    logger.info('Distance to target: %s: %.2f}' % (joint_type, feedback.feedback.distance_to_target))

class Client(Node):
    def __init__(self):
        super().__init__('joint_action_client')
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self._action_client = ActionClient(
            self, 
            JointTarget, 
            'joint_target')
        self.get_logger().info('Waiting for action server...')
        self._action_client.wait_for_server()

    def get_tf(self, j1, j2):
        time = self.get_clock().now()
        future = self.tf_buffer.wait_for_transform_async(j1, j2, time)
        rclpy.spin_until_future_complete(self, future)
        tf =  self.tf_buffer.lookup_transform(j1, j2, time)
        return tf

    def get_joint_distance(self, joint_type):
        if joint_type == "raise":
            trans = self.get_tf("arm_throttle", "cab")
            euler = tf_transformations.euler_from_quaternion([
                trans.transform.rotation.x,
                trans.transform.rotation.y,
                trans.transform.rotation.z,
                trans.transform.rotation.w     
            ])
            distance = euler[1]/pi*180

        elif joint_type == "extend":
            trans = self.get_tf("arm_end", "arm_throttle")
            distance = trans.transform.translation.x

        elif joint_type == "rotate":
            trans = self.get_tf("cab", "car")
            euler = tf_transformations.euler_from_quaternion([
                trans.transform.rotation.x,
                trans.transform.rotation.y,
                trans.transform.rotation.z,
                trans.transform.rotation.w     
            ])
            distance = euler[2]/pi*180

        elif joint_type == "lift":
            trans = self.get_tf("hook", "arm_end")
            distance = trans.transform.translation.z

        return distance

# 最基本的发出命令单元      
    def send_goal(self, order, joint_type):
        goal_msg = JointTarget.Goal()
        goal_msg.target = order
        goal_msg.joint_type = joint_type
        self.get_logger().info('Sending goal request...')
        send_goal_future = self._action_client.send_goal_async(
            goal_msg, feedback_callback=lambda feedback: feedback_cb(self.get_logger(), goal_msg.joint_type, feedback))
        rclpy.spin_until_future_complete(self, send_goal_future)
        goal_handle = send_goal_future.result()
        self.get_logger().info('Goal accepted :)')
        get_result_future = goal_handle.get_result_async()
        return goal_handle, get_result_future

# 指定各个关节的相对当前位置的位移位置（或相对位移），逐个发出命令.发送的请求值恒定为相对位移距离，如果设定为使用绝对距离，则先推算出相对位移距离。
    def operation_by_joint(self, _raise, _extend, _rotate, _lift, relative=True):
        self.cancel_flag = 0
        joint_type_list = ['raise', 'extend', 'rotate', 'lift']
        if not relative:
            now_raise = self.get_joint_distance('raise')
            absolute_raise = _raise - now_raise
            now_extend = self.get_joint_distance('extend')
            absolute_extend = _extend - now_extend
            now_rotate = self.get_joint_distance('rotate')
            absolute_rotate = _rotate - now_rotate
            now_lift = self.get_joint_distance('lift')
            absolute_lift = _lift - now_lift
            joint_move_list = [float(absolute_raise), float(absolute_extend), float(absolute_rotate), float(absolute_lift-absolute_extend/4)]                                    
        else:
            joint_move_list = [_raise, _extend, _rotate, _lift-_extend/4]

        self.get_logger().info('Target: '+str(joint_move_list))

        for joint_operation in (zip(joint_move_list,joint_type_list)): #逐个发送
            if self.cancel_flag == 1:
                break
            if int(abs(joint_operation[0])) <= 0.1: #Ignore too small distance operation
                continue
            self.goal_handle, self.get_result_future = self.send_goal(order = joint_operation[0], joint_type=joint_operation[1])
            try:
                rclpy.spin_until_future_complete(self, self.get_result_future)
                self.get_logger().info('Goal succeeded! {1}: Moved by {0}'.format(self.get_result_future.result().result.move_distance, joint_operation[1]))
            except:
                self.cancel_goal()
                return

# 后台运行，用于配合pyqt等多线程进程
    def operation_no_hup(self, _raise, _extend, _rotate, _lift, relative=True):
        try:
            if self.goal_handle:
                self.cancel_goal()
        except:
            pass
        t1 = threading.Thread(target = self.operation_by_joint, args=(_raise, _extend, _rotate, _lift, relative))
        t1.start()
        time.sleep(5)
        self.cancel_goal()

# 取消行动命令
    def cancel_goal(self):
        future = self.goal_handle.cancel_goal_async()
        self.cancel_flag = 1
        cancel_response = future.result()
        print("goal canceled")

    def get_hook_position(self):
        position = self.get_tf("car", "hook").transform.translation
        return position
    
    def operation_by_end(self, x_target, y_target, z_target, l_rt=1, l_e=1, l_l=1, l_rs=1, relative=False):
        now_end_position = self.get_hook_position()
        x_old = now_end_position.x
        y_old = now_end_position.y
        z_old = now_end_position.z
        theta_old = self.get_joint_distance("raise")/180*pi
        t1 = sqrt(x_target**2+y_target**2)
        t0 = sqrt(x_old**2+y_old**2)
        theta_target = root(lambda theta : l_e*(t1/cos(theta)-t0/cos(theta_old))*t1*sin(theta)/cos(theta)**2 + l_l*(tan(theta)*t1-z_target-tan(theta_old)*t0+z_old)*t1/cos(theta)**2+l_rs*(theta-theta_old), pi/2-0.1).x[0]
        raise_move = theta_target-theta_old
        rotate_move = atan(y_target/x_target) - atan(y_old/x_old)
        extend_move = t1/cos(theta_target) - t0/cos(theta_old)
        lift_move = t1*tan(theta_target) - t0*tan(theta_old) - z_target + z_old
        return (raise_move, extend_move, rotate_move, lift_move)

        

def main():
    if not rclpy.ok():
        rclpy.init()
    
    action_client = Client()

    #executor = MultiThreadedExecutor()
    #rclpy.spin(action_client, executor=executor)

    action_client.operation(5.0, 1.0, 0.0, 0.0)

if __name__ == '__main__':
    main()